{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5f1a1bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Forest_fire\n",
    "import numpy as np\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "from random import randint\n",
    "import copy\n",
    "import seaborn as sns\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "\n",
    "fig_num=0\n",
    "def random_start(size=[10,10],start_rate=0.8):\n",
    "    game_map=np.zeros(size)\n",
    "    for i in range(int(size[0]*size[1]*start_rate)):\n",
    "        game_map[np.random.randint(0,size[0]),np.random.randint(0,size[1])]=1\n",
    "    return game_map\n",
    "def output(game_map):\n",
    "    fig = plt.figure(figsize=(10,10))\n",
    "    #for y in range(len(game_map)):\n",
    "    #    for x in range(len(game_map[0])):\n",
    "    #        if game_map[y,x]==1:\n",
    "    #            plt.scatter(y,x,marker=\"s\",color=\"black\",s=200,)\n",
    "    #        if game_map[y,x]==0:\n",
    "    #            plt.scatter(y,x,marker=\"s\",color=\"white\",s=200)\n",
    "    ax = sns.heatmap(game_map,cmap=\"Greens\",cbar=False)      \n",
    "    return \"done\"\n",
    "def output_3d(game_map_arr):\n",
    "    fig=plt.figure(figsize=(10,10))\n",
    "    ax = fig.gca(projection='3d')\n",
    "    print(game_map_arr.shape)\n",
    "    for t in range(game_map_arr.shape[0]):\n",
    "        for y in range(game_map_arr.shape[1]):\n",
    "            for x in range(game_map_arr.shape[2]):\n",
    "                if game_map_arr[t,y,x]==1:\n",
    "                    ax.scatter(x,y,t,marker=\"s\",color=\"black\",alpha=0.1)\n",
    "    ax.set_xlabel('X')\n",
    "    ax.set_ylabel('Y')\n",
    "    ax.set_zlabel('t')\n",
    "    #plt.savefig(\"D:/3_bodies/pic_orbit/\"+str(save_i)+\".jpg\")\n",
    "            #img=cv2.imread(path+str(i)+\".jpg\")\n",
    "            #post_video.write(plt)\n",
    "    plt.show()\n",
    "class forest(object):\n",
    "    def __init__(self,game_map,size=[10,10],fig_num=0,rule=[3,0,2,0.03,0.01]):\n",
    "        self.rule=rule#[r,a,t,p]\n",
    "        self.size=size\n",
    "        self.fig_num=fig_num\n",
    "        self.game_map=copy.deepcopy(game_map)\n",
    "        self.time_map=copy.deepcopy(game_map)*self.rule[2]#t\n",
    "        \n",
    "    \n",
    "    def main(self,time=100,map_out=False,pic_out=True,save_fig=False,save_path=None):   \n",
    "        reward_line=[]\n",
    "        game_map_arr=[]\n",
    "        for i in range(time):\n",
    "    #plt.colorbar(game_map,shrink=0.8,autoscale_None=None)   \n",
    "            if map_out and i%1==0:\n",
    "                output(self.game_map)\n",
    "                if save_fig:\n",
    "                    plt.savefig(save_path+str(self.fig_num)+\"-\"+str(i)+\".png\")\n",
    "                os.system(\"cls\")\n",
    "                plt.show()\n",
    "            old_map=copy.deepcopy(self.game_map)\n",
    "            self.refresh()\n",
    "            game_map_arr.append(self.game_map)\n",
    "            new_map=self.game_map\n",
    "            reward_line.append(self.reward_func(old_map,new_map))\n",
    "        game_map_arr=np.array(game_map_arr)\n",
    "        #output_3d(game_map_arr)\n",
    "        if pic_out:\n",
    "            plt.plot(reward_line)    \n",
    "            plt.show()\n",
    "        return sum(reward_line[-30:])/len(reward_line[-30:])\n",
    "    def reward_func(self,old_map,new_map):\n",
    "        delta_map=np.abs(new_map-old_map)\n",
    "        reward=np.sum(delta_map)/(len(new_map)*len(new_map[0]))\n",
    "        return reward\n",
    "    def near_num(self,point):\n",
    "        num=0\n",
    "        around=range(-self.rule[0],self.rule[0]+1,1)\n",
    "        for i in around:\n",
    "            for j in around:\n",
    "                if point[0]+i>=0 and point[0]+i<len(self.game_map)-1 and point[1]+j>=0 and point[1]+j<len(self.game_map[0])-1:\n",
    "                    if self.game_map[point[0]+i][point[1]+j]==0.5 and (i!=0 or j!=0):\n",
    "                        num+=1\n",
    "        return num\n",
    "    def refresh(self):\n",
    "        new_map=copy.deepcopy(self.game_map)\n",
    "        for x in range(self.game_map.shape[0]):\n",
    "            for y in range(self.game_map.shape[1]):\n",
    "                if self.time_map[x,y]>0:\n",
    "                    self.time_map[x,y]-=1\n",
    "                if self.time_map[x,y]==1 and self.game_map[x,y]==0.5:\n",
    "                    new_map[x,y]=0\n",
    "                \n",
    "                num=self.near_num([x,y])\n",
    "                if num>self.rule[1] and self.game_map[x,y]==1:#a\n",
    "                    new_map[x][y]=0.5\n",
    "                    self.time_map[x,y]=self.rule[2]\n",
    "                elif np.random.rand()<self.rule[3] and self.game_map[x,y]==1:\n",
    "                    new_map[x][y]=0.5\n",
    "                    self.time_map[x,y]=self.rule[2]\n",
    "                elif np.random.rand()<self.rule[4] and self.game_map[x,y]==0:\n",
    "                    new_map[x][y]=1\n",
    "        self.game_map=new_map\n",
    "        return 0\n",
    "def ran_start(size=[10,10],time=10):\n",
    "    r_max=0\n",
    "    for i in range(time):\n",
    "        start_map=random_start(size=size)\n",
    "        live_game_1=live_game(start_map)\n",
    "        r=live_game_1.main(pic_out=False)\n",
    "        print(r)\n",
    "        if r>r_max:\n",
    "            out_map=copy.deepcopy(start_map)\n",
    "            r_max=copy.deepcopy(r)\n",
    "    return out_map\n",
    "class random_walk(object):\n",
    "    def __init__(self,start_map,epoch=100):\n",
    "        self.epoch=10\n",
    "        self.start_map=copy.deepcopy(start_map)\n",
    "        self.r_max=0\n",
    "        self.game_map=copy.deepcopy(start_map)\n",
    "    def step(self):\n",
    "        ran_x,ran_y=int(np.random.rand()*len(self.game_map)),int(np.random.rand()*len(self.game_map[0]))\n",
    "        #print(ran_x,ran_y))\n",
    "        self.game_map[ran_x,ran_y]=np.abs(self.game_map[ran_x,ran_y]-1)\n",
    "        live_game_1=live_game(self.game_map)\n",
    "        r=live_game_1.main(pic_out=False)\n",
    "        if r>self.r_max:\n",
    "            self.start_map=self.game_map\n",
    "            self.r_max=r\n",
    "        if r>0.1:\n",
    "            return True\n",
    "        else:\n",
    "            return False\n",
    "    def main(self,episode=200):\n",
    "        for i_episode in range(episode):\n",
    "            live_game_1=live_game(self.start_map)\n",
    "            self.r_max=live_game_1.main(pic_out=False)\n",
    "            for i_epoch in range(self.epoch):\n",
    "                if self.step():\n",
    "                    print(\"find!r=\",self.r_max)\n",
    "                    output(self.game_map)\n",
    "                    plt.show()\n",
    "                    return self.game_map    \n",
    "            print(\"EP\"+str(i_episode)+\":r_max=\"+str(self.r_max))\n",
    "            self.game_map=self.start_map\n",
    "        output(self.game_map)\n",
    "        plt.show()\n",
    "        return None\n",
    "game_act=[]\n",
    "save_path=r\"D:/3_bodies/pic/Fire_\"\n",
    "start_map=random_start(size=[200,200],start_rate=0.9)\n",
    "forest_0=forest(start_map,fig_num=\"test\",rule=[1,0,2,1e-4,1e-3])\n",
    "r=forest_0.main(time=1000,pic_out=True,map_out=True,save_fig=True,save_path=save_path)\n",
    "print(r)\n",
    "\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12ad39d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#imgs2video\n",
    "import os\n",
    "import cv2\n",
    "from PIL import Image\n",
    "fourcc = cv2.VideoWriter_fourcc('X', 'V', 'I', 'D') \n",
    "path=r\"D:/3_bodies/pic/Fire_test-\"#input(\"path?\")\n",
    "save_path=r\"D:/3_bodies/fire_g_e-3.mp4\"#input(\"save_path?\")\n",
    "post_video=cv2.VideoWriter(save_path, fourcc, 30,(720,720))\n",
    "for i in range(1000):\n",
    "    print(path+str(i)+\".png\")\n",
    "    img=cv2.imread(path+str(i)+\".png\")\n",
    "    post_video.write(img)\n",
    "    del img\n",
    "    #if filename==\"0.jpg\":\n",
    "        #Image.fromarray(img).show()\n",
    "\n",
    "post_video.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "083074aa",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
